{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wlFbFLUghfjo"
      },
      "source": [
        "##### Copyright 2022 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "4FyfuZX-gTKS"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I_16rv9EPhB_"
      },
      "source": [
        "# Passage Ranking using TFR-BERT\n",
        "\n",
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/ranking/tutorials/tfr_bert\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/ranking/blob/master/docs/tutorials/tfr_bert.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/ranking/blob/master/docs/tutorials/tfr_bert.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/ranking/docs/tutorials/tfr_bert.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "V8tMYn22vtDV"
      },
      "source": [
        "TensorFlow Ranking can handle heterogeneous dense and sparse features, and scales up to millions of data points. However, building and deploying a learning to rank model to operate at scale creates additional challenges beyond simply designing a model. The Ranking library provides workflow utility classes for building [distributed training](https://www.tensorflow.org/guide/distributed_training) for large-scale ranking applications. For more information about these features, see the TensorFlow Ranking [Overview](../overview).\n",
        "\n",
        "This tutorial shows you how to build a ranking model that uses BERT for scoring. [BERT](https://github.com/google-research/bert) is a highly effective pretrained module to effective encode textual features into contextualized word embeddings. We use BERT to initialize the ranking model and finetune the model using a ranking loss.\n",
        "\n",
        "Note: An advanced version of this code is also available as a [Python script](https://github.com/tensorflow/ranking/blob/master/tensorflow_ranking/examples/keras/tfrbert_antique_train.py)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UxG7i8xbDIDF"
      },
      "source": [
        "## ANTIQUE dataset\n",
        "\n",
        "In this tutorial, you will build a ranking model for ANTIQUE, a question-answering dataset using BERT as the scoring function. Bidirectional Encoder Representations from Transformers (BERT) is a transformer-based machine learning technique which has proven to be effective in many natural language processing (NLP) tasks. Recent work on [TFR-BERT](https://https://arxiv.org/abs/2004.08476) has shown BERT to be an effective scoring function for learning-to-rank tasks.\n",
        "\n",
        "Given a query, and a list of answers, the objective of the ranking model is to rank the answers with optimal rank related metrics, such as NDCG. For more details about ranking metrics, review evaluation measures [offline metrics](https://en.wikipedia.org/wiki/Evaluation_measures_(information_retrieval)#Offline_metrics).\n",
        "\n",
        "[ANTIQUE](https://ciir.cs.umass.edu/downloads/Antique/) is a publicly available dataset for open-domain non-factoid question answering, collected from Yahoo! answers.\n",
        "Each question has a list of answers, whose relevance are graded on a scale of 0-4, 0 for irrelevant and 4 for fully relevant.\n",
        "The list size can vary depending on the query, so we use a fixed \"list size\" of 50, where the list is either truncated or padded with default values.\n",
        "The dataset is split into 2206 queries for training and 200 queries for testing. For more details, please read the technical paper on [arXiv](https://arxiv.org/abs/1905.08957)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ucWaXnFazZXD"
      },
      "source": [
        "## Setup\n",
        "\n",
        "Download and install the TensorFlow Ranking and TensorFlow Model Garden packages."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "aPmhLkMWgPLO"
      },
      "outputs": [],
      "source": [
        "!pip install -q tensorflow-ranking tf-models-official"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9OKDJUjq0rnm"
      },
      "source": [
        "Import TensorFlow Ranking and useful libraries through the notebook."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fmlaz2D5Ux3J"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "import tensorflow as tf\n",
        "import tensorflow_ranking as tfr\n",
        "from official.nlp.configs import encoders\n",
        "from tensorflow_ranking.extension.premade import tfrbert_task"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JNilCoqq1jJn"
      },
      "source": [
        "## Data preparation\n",
        "\n",
        "Download training and test data."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Mwxtsi4wqoOJ"
      },
      "outputs": [],
      "source": [
        "!wget -O \"/tmp/train.tfrecords\" \"https://ciir.cs.umass.edu/downloads/Antique/tf-ranking/antique_train_seq_64_elwc.tfrecords\"\n",
        "!wget -O \"/tmp/test.tfrecords\" \"https://ciir.cs.umass.edu/downloads/Antique/tf-ranking/antique_test_seq_64_elwc.tfrecords\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D5vL2R7aOSoe"
      },
      "outputs": [],
      "source": [
        "!mkdir -p /tmp/tfrbert\n",
        "!wget \"https://storage.googleapis.com/cloud-tpu-checkpoints/bert/v3/uncased_L-12_H-768_A-12.tar.gz\" -P \"/tmp/tfrbert\"\n",
        "!mkdir -p /tmp/tfrbert/uncased_L-12_H-768_A-12\n",
        "!tar -xvf /tmp/tfrbert/uncased_L-12_H-768_A-12.tar.gz --strip-components 3 -C \"/tmp/tfrbert/uncased_L-12_H-768_A-12/\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tFbFBTUh9WXf"
      },
      "source": [
        "## Overview of TFR-BERT in Orbit\n",
        "\n",
        "BERT-based ranking models ([TFR-BERT](https://arxiv.org/abs/2004.08476)) have been shown to be effective for learning-to-rank tasks when using raw textual features for query and passages in MSMARCO passage ranking dataset.\n",
        "\n",
        "[Orbit](https://github.com/tensorflow/models/tree/master/orbit) is a flexible, lightweight library designed to make it easy to write custom training loops in TensorFlow. TensorFlow Ranking provides support for implementing ranking models, particularly for BERT based ranking models using Orbit."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "aQ-VTA56sOTA"
      },
      "source": [
        "## Create a Ranking Task for TFR-BERT\n",
        "\n",
        "We create a ranking task for TFR-BERT model which can be trained using Orbit. The steps to build this are:\n",
        "\n",
        "1.   Define Feature Specifications\n",
        "2.   Define datasets\n",
        "5.   Setup data and task configurations\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "at0nVKnts8Pn"
      },
      "source": [
        "### Specify Features\n",
        "\n",
        "[Feature Specification](https://www.tensorflow.org/api_docs/python/tf/io) are TensorFlow abstractions to capture information about each feature. These help developers and model researchers understand and use a model.\n",
        "\n",
        "Create feature specifications for context features, example features, and labels, consistent with the input formats for ranking, such as ELWC format."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nSXd4pEPqaQW"
      },
      "outputs": [],
      "source": [
        "SEQ_LENGTH = 64\n",
        "context_feature_spec = {}\n",
        "example_feature_spec = {\n",
        "    'input_word_ids': tf.io.FixedLenFeature(\n",
        "        shape=(SEQ_LENGTH,), dtype=tf.int64,\n",
        "        default_value=[0] * SEQ_LENGTH),\n",
        "    'input_mask': tf.io.FixedLenFeature(\n",
        "        shape=(SEQ_LENGTH,), dtype=tf.int64,\n",
        "        default_value=[0] * SEQ_LENGTH),\n",
        "    'input_type_ids': tf.io.FixedLenFeature(\n",
        "        shape=(SEQ_LENGTH,), dtype=tf.int64,\n",
        "        default_value=[0] * SEQ_LENGTH)}\n",
        "label_spec = (\n",
        "    \"relevance\",\n",
        "    tf.io.FixedLenFeature(shape=(1,), dtype=tf.int64, default_value=-1)\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X2Iq_YA2HpCe"
      },
      "source": [
        "Note: the `default_value` of `label_spec` feature is set to -1 to take care of the padding items to be masked out."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nKoHvr1oj02f"
      },
      "source": [
        "### Define Datasets\n",
        "\n",
        "We define data configurations for training and validation data, which specifies parameters such as path, batch size, and dataset format. These configurations are used to create training and validation datasets."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fCJZaAqBj0YE"
      },
      "outputs": [],
      "source": [
        "# Set up data config\n",
        "# We use a small list size here for demo purposes only. Users can use a larger\n",
        "# list size on a machine with more memory to train TFR-BERT.\n",
        "train_data_config = tfrbert_task.TFRBertDataConfig(\n",
        "    input_path=\"/tmp/train.tfrecords\",\n",
        "    is_training=True,\n",
        "    global_batch_size=8,\n",
        "    list_size=2,\n",
        "    dataset_fn='tfrecord',\n",
        "    seq_length=64)\n",
        "\n",
        "validation_data_config = tfrbert_task.TFRBertDataConfig(\n",
        "    input_path=\"/tmp/test.tfrecords\",\n",
        "    is_training=False,\n",
        "    global_batch_size=8,\n",
        "    list_size=2,\n",
        "    dataset_fn='tfrecord',\n",
        "    seq_length=64)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8-O_z1S2j_s0"
      },
      "source": [
        "### Define Task\n",
        "\n",
        "We define a task configuration which defines the training and validation dataset along with the model. This configuration creates a `TFRBertTask` object that can be trained using Orbit."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uLly4jcykCC1"
      },
      "outputs": [],
      "source": [
        "# Set up task config\n",
        "task_config = tfrbert_task.TFRBertConfig(\n",
        "    init_checkpoint='/tmp/tfrbert/uncased_L-12_H-768_A-12/bert_model.ckpt',\n",
        "    train_data=train_data_config,\n",
        "    validation_data=validation_data_config,\n",
        "    model=tfrbert_task.TFRBertModelConfig(\n",
        "        encoder=encoders.EncoderConfig(\n",
        "            bert=encoders.BertEncoderConfig(num_layers=12))))\n",
        "\n",
        "# Set up TFRBertTask\n",
        "task = tfrbert_task.TFRBertTask(\n",
        "    task_config,\n",
        "    label_spec=label_spec,\n",
        "    dataset_fn=tf.data.TFRecordDataset,\n",
        "    logging_dir='/tmp/model_dir')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_EPvXEbomK29"
      },
      "source": [
        "## Train and evaluate the model\n",
        "\n",
        "We define the training loop here to train and evaluate the model. We define the metrics, create train and eval datasets and train the model for a specific number of training steps.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BZr8MX6VmQSj"
      },
      "outputs": [],
      "source": [
        "metrics = task.build_metrics()\n",
        "model = task.build_model()\n",
        "task.initialize(model)\n",
        "train_dataset = task.build_inputs(task_config.train_data)\n",
        "vali_dataset = task.build_inputs(task_config.validation_data)\n",
        "train_iterator = iter(train_dataset)\n",
        "vali_iterator = iter(vali_dataset)\n",
        "optimizer = tf.keras.optimizers.Adam(lr=1e-6)\n",
        "\n",
        "NUM_TRAIN_STEPS = 100\n",
        "EVAL_STEPS = 10\n",
        "for train_step in range(NUM_TRAIN_STEPS):\n",
        "  task.train_step(next(train_iterator), model, optimizer, metrics=metrics)\n",
        "  train_metrics = {m.name: m.result().numpy() for m in metrics}\n",
        "  print(\"Training metrics for epoch: \" + str(train_step) + \" \", train_metrics)\n",
        "\n",
        "  if train_step % EVAL_STEPS == 0:\n",
        "    task.validation_step(next(train_iterator), model, metrics=metrics)\n",
        "    vali_metrics = {m.name: m.result().numpy() for m in metrics}\n",
        "    print(\"Validation metrics for epoch: \" + str(train_step) + \" \",\n",
        "          vali_metrics)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "tfr_bert_orbit.ipynb",
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
